// base.js
// base proxy class
const fs = require('fs');
const path = require('path');
const { CoCoMetricDB } = require('../metrics/CoCoMetricDB');
const { getAllFilesInDirectory } = require('./fileUtils');

class BaseProxy {
    constructor() { }

    // 返回训练数据可视化配置
    // 主要是对于scalars中包含的key的展示配置
    visConfig(train_cfg) {
        // console.log('train_cfg = ', train_cfg);
        train_cfg = typeof train_cfg === 'string' ? JSON.parse(train_cfg) : train_cfg;
        let config = {
            charts: ['loss', 'acc', 'lr', 'mem'],
            chart: {
                loss: {
                    title: 'Loss',
                    keys: ['loss', 'loss_rpn_cls', 'loss_rpn_bbox', 'loss_cls', 'loss_bbox']
                },
                acc: {
                    title: 'Accuracy',
                    keys: ['acc']
                },
                lr: {
                    title: 'Learning Rate',
                    keys: ['lr']
                },
                mem: {
                    title: 'Memory',
                    keys: ['memory']
                }
            },
            evaluation: {
                keys: []
            }
        };
        if (train_cfg && train_cfg.dataset_type == 'voc') {
            config.evaluation.keys = ['pascal_voc/mAP', 'pascal_voc/AP50'];
        } else if (train_cfg && train_cfg.dataset_type == 'coco') {
            config.evaluation.keys = ['coco/bbox_mAP', 'coco/bbox_mAP_50', 'coco/bbox_mAP_75', 'coco/bbox_mAP_s', 'coco/bbox_mAP_m', 'coco/bbox_mAP_l'];
        }
        return config;
    }

    // 训练过程中，从日志中提取出数据并返回用于展示
    loadRuntimeLogs(mode, logPath, data, train_cfg) {
        try {
            const fileName = path.basename(logPath).toLowerCase();
            const ext = path.extname(logPath).toLowerCase();
            if (ext === '.log') {
                return {
                    event: 'stdio',
                    params: {
                        fileName,
                        logData: data
                    }
                };
            } else if (ext === '.json') {
                // TODO: 看起来这里会有读取数据不是一个合格json的情况
                if (fileName === 'scalars.json') {
                    return {
                        event: 'scalars',
                        params: {
                            fileName,
                            logData: data
                        }
                    };
                } else {
                    return null;
                }
            } else if (ext === '.py') {
                if (fileName === 'config.py') {
                    return {
                        event: 'config',
                        params: {
                            fileName,
                            logData: data
                        }
                    };
                }
            }
        } catch (error) {
            console.error(`Error handling log file: `, error);
            return null;
        }
    }

    // 训练结束后，从本地文件提取数据并返回用于展示
    async loadTrainResults(output_dir, train_cfg) {
        let modelFiles = await getAllFilesInDirectory(output_dir, ['.pth'], []);
        let logFiles = await getAllFilesInDirectory(output_dir, ['.json', '.py', '.log'], ['test']);
        let logPath, scalarsPath, evalsPath, configPath;
        for (let i = 0; i < logFiles.length; i++) {
            let filePath = logFiles[i];
            let basename = path.basename(filePath);
            let extname = path.extname(filePath);
            if (extname == '.log') {
                logPath = filePath;
            } else if (basename == 'config.py') {
                configPath = filePath;
            } else if (basename == 'scalars.json') {
                let paths = this.splitScalars(output_dir, filePath, modelFiles, train_cfg);
                scalarsPath = paths.scalarsPath;
                evalsPath = paths.evalsPath;
            }
        }

        let result = {
            logPath,
            scalarsPath,
            evalsPath,
            configPath,
            modelFiles
        };

        return result;
    }

    async getEvalResultList(output_dir, train_cfg) {
        train_cfg = typeof train_cfg === 'string' ? JSON.parse(train_cfg) : train_cfg;
        let modelFiles = await getAllFilesInDirectory(output_dir, ['.pth'], []);
        let eval_result_list = [];
        for (let i = 0; i < modelFiles.length; i++) {
            let model_path = modelFiles[i];
            let modelName = path.basename(model_path, '.pth');
            let eval_results_dir = path.join(output_dir, 'test', `test_${modelName}`);
            console.log(eval_results_dir);

            let eval_result = {};
            if (fs.existsSync(eval_results_dir)) {
                eval_result.model_path = model_path;
                eval_result.existed = true;
            } else {
                eval_result.model_path = model_path;
                eval_result.existed = false;
            }
            eval_result_list.push(eval_result);
        }
        return eval_result_list;
    }

    // 评测结束后，从本地读取评测结果和详情信息
    async loadEvalResults(sqlite3, output_dir, model_path, filter, train_cfg) {
        train_cfg = typeof train_cfg === 'string' ? JSON.parse(train_cfg) : train_cfg;

        // console.log(`loadEvalResults output_dir = ${output_dir}, model_path = ${model_path}, filter = `, filter);
        const modelName = path.basename(model_path, '.pth');

        // 返回评测的整体结果
        let eval_results = {};
        const eval_results_path = path.join(output_dir, `test_${modelName}`, 'eval_results.json');
        // console.log('eval_result_path = ', eval_results_path);
        try {
            let jsonData = fs.readFileSync(eval_results_path, 'utf8');
            eval_results = JSON.parse(jsonData);
            // console.log(eval_results);
        } catch (error) {
            // console.error('eval_results.json not found!');
            return null;
        }

        // 返回具体的data_samples数据
        let data_samples = [];
        const data_samples_path = path.join(output_dir, `test_${modelName}`, 'data_samples.json');
        // console.log('data_samples_path = ', data_samples_path);
        try {
            let jsonData = fs.readFileSync(data_samples_path, 'utf8');
            jsonData = JSON.parse(jsonData);
            data_samples = jsonData.map((item) => {
                return {
                    ...item,
                    img_name: path.basename(item.img_path)
                };
            });
            // console.log(eval_results);
        } catch (error) {
            // console.error('data_samples.json not found!');
            return null;
        }

        // 从数据库根据filter读取对应的详情结果
        let eval_details = [];
        try {
            if (train_cfg && train_cfg.dataset_type == 'voc') {
                console.log('voc eval_details.db not implemented!');
            } else if (train_cfg && train_cfg.dataset_type == 'coco') {
                let metric_db_path = path.join(output_dir, `test_${modelName}`, 'eval_details.db');
                let metric_db = new CoCoMetricDB(sqlite3, metric_db_path);
                eval_details = await metric_db.getEvalDetails(filter);
                await metric_db.close();
            }
        } catch (error) {
            console.error('eval_details.db not found!');
            return null;
        }

        // 返回结果
        return {
            data_samples,
            eval_results,
            eval_details
        };
    }

    // 一些实用函数
    // 单独处理scalars.json文件，区分lr和评测结果
    splitScalars(output_dir, jsonPath, modelFiles, train_cfg) {
        const jsonData = fs.readFileSync(jsonPath, 'utf8');
        const lines = jsonData.split('\n');
        // 读取原始的scalars.json文件
        let scalars = [],
            evals = [];
        for (let i = 0; i < lines.length; i++) {
            let line = lines[i].trim();
            if (line.length == 0) {
                continue;
            }

            try {
                const jsonLine = JSON.parse(line);
                if ('lr' in jsonLine) {
                    scalars.push(jsonLine);
                } else {
                    evals.push(jsonLine);
                }
            } catch (error) {
                console.error(error);
            }
        }

        // 处理评测结果的过滤
        let config = this.visConfig(train_cfg);
        // console.log('config = ', config);

        let filteredData = evals.map((item) => {
            let modelPath = modelFiles.find((file) => file.includes(`epoch_${item.step}.pth`));
            let metrics = config.evaluation.keys.reduce((acc, key) => {
                acc[key] = item[key];
                return acc;
            }, {});
            // console.log('metrics = ', metrics);
            // 判断是否存在评测结果
            let modelName = path.basename(modelPath, '.pth');
            let eval_results_dir = path.join(output_dir, 'test', `test_${modelName}`);
            let exist_eval = false;
            if (fs.existsSync(eval_results_dir)) {
                exist_eval = true;
            }
            return {
                step: item['step'],
                model_path: modelPath,
                exist_eval: exist_eval,
                ...metrics
            };
        });

        // 获取目录路径
        const dirPath = path.dirname(jsonPath);
        const scalarsPath = path.join(dirPath, 'scalars_.json');
        fs.writeFileSync(scalarsPath, JSON.stringify(scalars, null, 2), 'utf8');
        const evalsPath = path.join(dirPath, 'evals_.json');
        fs.writeFileSync(evalsPath, JSON.stringify(filteredData, null, 2), 'utf8');

        // 返回新文件的路径
        return { scalarsPath, evalsPath };
    }

    // 读取模型的配置文件列表
    loadConfigList() {
        throw new Error("Handler method 'loadConfigList' must be implemented by subclass.");
    }

    // 读取模型结构
    loadModelArch(config_name) {
        throw new Error("Handler method 'loadModelArch' must be implemented by subclass.");
    }

    // 读取模型参数
    loadModelParams(config_name) {
        throw new Error("Handler method 'loadModelParams' must be implemented by subclass.");
    }
}

module.exports = { BaseProxy };
